#include <stdio.h>
#include <cuda.h>
#include <sys/time.h>
#include <cuda_runtime.h>
#include <stdint.h>
#include "cublas_v2.h"

#define M 2048
#define N 2048
#define K 2048


# define timersub(a, b, result)                                         \
	do {                                                            \
		(result)->tv_sec = (a)->tv_sec - (b)->tv_sec;           \
		(result)->tv_usec = (a)->tv_usec - (b)->tv_usec;        \
		if ((result)->tv_usec < 0) {                            \
			--(result)->tv_sec;                             \
			(result)->tv_usec += 1000000;                   \
		}                                                       \
	} while (0)

static texture<float,1>  tex_x_float_A;
static texture<float,1>  tex_x_float_B;

static __inline__ __device__ float fetch_x_A(const int& i)
{
	return tex1Dfetch(tex_x_float_A, i);
}

static __inline__ __device__ float fetch_x_B(const int& i)
{
	return tex1Dfetch(tex_x_float_B, i);
}
	__global__ void 
kernel6(float *C, const float *A, const float *B, int m, int n,  int k,  int lda,  int ldb, int ldc, int offsetA, int offsetB) 
{
	const  int tx = threadIdx.x;
	const  int ty = threadIdx.y;

	const int iby = blockIdx.y * 96;
	const int ibx = blockIdx.x * 96;
	const int idt = ty * 64 + tx;

	const int tx2 = idt%16;
	const int ty2 = idt/16;

	__shared__ float Bb[16][97];
	__shared__ float Abs[96][17];

	float xxA[6];
	float xxB[6];

	int trackA = offsetA + __mul24( ibx + ty2, lda) + tx2;
	int trackB = offsetB + iby+ tx2 + __mul24(ty2, ldb);

	A += trackA; 
	B += trackB; 

	int tll = tx2; 

#pragma unroll
	for(int y=0; y<6; y++)
		Abs[ty2+16*y][tx2] = fetch_x_A(trackA +  lda*16*y);

#pragma unroll
	for(int y=0; y<6; y++)
		Bb[ty2][tx2+16*y] = fetch_x_B(trackB+16*y);

	__syncthreads();

	const float *Bend = B + k*ldb - 16*ldb;

	float Axs[6];
	float Bxp[6];

	float Cb[36] = {0,0,0,0,0,0, 0,0,0,0,0,0,  0,0,0,0,0,0, 0,0,0,0,0,0,
		0,0,0,0,0,0, 0,0,0,0,0,0};
	do 
	{
		tll+=16;
		A += 16;
		B += 16*ldb;
		trackA+=16; 
		trackB+=16*ldb;

#pragma unroll
		for( int y=0; y<6; y++)
			xxA[y] = fetch_x_A(trackA + lda*y*16);

#pragma unroll
		for( int y=0; y<6; y++)
			xxB[y] = fetch_x_B(trackB + 16*y);

#pragma unroll 
		for( int j1=0;j1<16;j1++)
		{
#pragma unroll
			for( int y=0; y<6; y++)
				Axs[y] =  Abs[tx2 + y*16][j1];

#pragma unroll
			for( int y=0; y<6; y++)
				Bxp[y]= Bb[j1][ty2 + y*16];

#pragma unroll 
			for( int x=0; x<6; x++)
#pragma unroll 
				for( int y=0;y<6;y++)
					Cb[x*6+y] += Axs[x]*Bxp[y];
		}
		__syncthreads();

#pragma unroll
		for( int y=0; y<6; y++)
			Abs[ty2 + 16*y][tx2] = xxA[y];

#pragma unroll
		for( int y=0; y<6; y++)
			Bb[ty2][tx2+y*16] = xxB[y];

		__syncthreads();
	} 
	while (B < Bend);

	C += tx2 + ibx  + __mul24 (ty2 +  iby ,ldc);

#pragma unroll 
	for( int j1=0; j1<16; j1++)
	{
#pragma unroll
		for( int y=0; y<6; y++)
			Axs[y] =  Abs[tx2 + y*16][j1];

#pragma unroll
		for( int y=0; y<6; y++)
			Bxp[y]= Bb[j1][ty2 + y*16];

#pragma unroll 
		for( int x=0; x<6; x++)
#pragma unroll 
			for( int y=0; y<6; y++)
				Cb[x*6+y]  += Axs[x]*Bxp[y];
	}

	int gy = iby + ty2;
#pragma unroll
	for( int y=0; y<6; y++, gy+=16)
	{
		int gx = ibx + tx2; 
#pragma unroll
		for(int x=0; x<6; x++, gx+=16)
		{
			if (gx < m && gy < n)
				C[x*16] = Cb[y+x*6] + 0 * C[x*16];
		}

		C+=ldc*16;
	}

}

int main(int argc, const char *argv[])
{
	struct timeval time_start, time_end, time_res;

	float *m1_h, *m2_h, *o_h;
	float *m1_d, *m2_d, *o_d;
	float *o_h2;

	float alpha = 1.0;
	float beta = 0.0;

	int count = 0;

	m1_h = (float *) malloc(M * K * sizeof(float));
	m2_h = (float *) malloc(K * N * sizeof(float));
	o_h = (float *) malloc(M * N * sizeof(float));
	o_h2 = (float *) malloc(M * N * sizeof(float));
	

	cudaMalloc( (void **)&m1_d, M * K * sizeof(float));
	cudaMalloc( (void **)&m2_d, K * N * sizeof(float));
	cudaMalloc( (void **)&o_d, M * N * sizeof(float));
	
	// Initialize host array and copy it to CUDA device
	for (int r = 0 ; r < M; ++r) {
		for (int c = 0 ; c < K; ++c) {
			m1_h[r*K + c] = rand()/(float)RAND_MAX;
		}
	}

	for (int r = 0 ; r < K; ++r) {
		for (int c = 0 ; c < N; ++c) {
			m2_h[r*N + c] = rand()/(float)RAND_MAX;
		}
	}

	cudaMemcpy(m1_d, m1_h, M * K * sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(m2_d, m2_h, K * N * sizeof(float), cudaMemcpyHostToDevice);

	size_t sizeA = K * M;
	size_t sizeB = N * K;

	size_t offsetA = 0;
	size_t offsetB = 0;
	
	cublasHandle_t handle;

	if(cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS){
		printf("cublasCteate fail\n");
		exit(-1);
	}

	if (cudaBindTexture(&offsetA, tex_x_float_A, m1_d, sizeA * sizeof(float)) != cudaSuccess) {
		printf("can not bind to texture\n");
		exit(-1);
	}
	if (cudaBindTexture(&offsetB, tex_x_float_B, m2_d, sizeB * sizeof(float)) != cudaSuccess) {
		printf("can not bind to texture\n");
		exit(-1);
	}

	dim3 threads( 64, 4 );
	dim3 grid( (M + 95)/96, (N + 95)/96 );

	cudaDeviceSynchronize();

	gettimeofday(&time_start, NULL);

	kernel6 <<< grid, threads >>> ( o_d, m1_d, m2_d, M, N, K , K, N, M ,offsetA, offsetB);
	cudaDeviceSynchronize();

	gettimeofday(&time_end, NULL);

	timersub(&time_end, &time_start, &time_res);

	printf("%ld,%03ld ms\n", time_res.tv_sec, time_res.tv_usec/1000);

	cudaUnbindTexture (tex_x_float_A);
	cudaUnbindTexture (tex_x_float_B);

	cudaMemcpy(o_h,o_d,M*N*sizeof(float), cudaMemcpyDeviceToHost);
	cudaDeviceSynchronize();

	gettimeofday(&time_start,NULL);

	cublasSgemm(handle, CUBLAS_OP_T,CUBLAS_OP_T, M, K, N, &alpha, m1_d, K, m2_d, N, &beta, o_d, N);
	cudaDeviceSynchronize();

	gettimeofday(&time_end,NULL);

	timersub(&time_end, &time_start, &time_res);

	printf("%ld,%03ld ms\n", time_res.tv_sec, time_res.tv_usec/1000);
	
	cudaMemcpy(o_h2, o_d, M*N*sizeof(float), cudaMemcpyDeviceToHost);

	for(int i = 0 ; i < M ; i ++){
		for(int j = 0 ; j < N ; j++){
			if( (o_h[i*N+j] - o_h2[i*N+j]) > 0.001 || (o_h[i*N+j] - o_h2[i*N+j]) < -0.001){
				count++;
			}
		}
	}
	printf("number of differences : %d\n",count);

	free(m1_h);    
	free(m2_h);
	free(o_h);
	free(o_h2);

	cudaFree(m1_d);
	cudaFree(m2_d);
	cudaFree(o_d);
}

